#!/usr/bin/env python3
# -*- coding: utf-8 -*-

#import json
import logging
import os
import time
import tensorflow as tf

import utils
from model import Model
from utils import read_data

from flags import parse_args
FLAGS, unparsed = parse_args()


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s', level=logging.DEBUG)


vocabulary = read_data(FLAGS.text)
print('Data size', len(vocabulary))

'''
with open(FLAGS.dictionary, encoding='utf-8') as inf:
    dictionary = json.load(inf, encoding='utf-8')

with open(FLAGS.reverse_dictionary, encoding='utf-8') as inf:
    reverse_dictionary = json.load(inf, encoding='utf-8')
'''

model = Model(learning_rate=FLAGS.learning_rate, batch_size=FLAGS.batch_size, num_steps=FLAGS.num_steps)
model.build()


with tf.Session() as sess:
    summary_string_writer = tf.summary.FileWriter(FLAGS.output_dir, sess.graph)

    saver = tf.train.Saver(max_to_keep=5)
    sess.run(tf.global_variables_initializer())
    sess.run(tf.local_variables_initializer())
    logging.debug('Initialized')

    try:
        checkpoint_path = tf.train.latest_checkpoint(FLAGS.output_dir)
        saver.restore(sess, checkpoint_path)
        logging.debug('restore from [{0}]'.format(checkpoint_path))

    except Exception:
        logging.debug('no check point found....')

    # Train network
    step = 0
    #对各种状态进行初始化
    new_state = sess.run(model.state_tensor)
    for x, y in utils.get_train_data(vocabulary, batch_size=FLAGS.batch_size, num_steps=FLAGS.num_steps):
        step += 1
        start = time.time()
        feed = {model.X: x,
                model.Y: y,
                model.keep_prob: 0.5,
                model.state_tensor: new_state}
        batch_loss, new_state, _ = sess.run([model.loss,
                                             model.final_state,
                                             model.optimizer],
                                             feed_dict=feed)

        end = time.time()
        # control the print lines
        max_steps=100000
        save_path='/home/paul/桌面/写诗器-homework.3/code'
        if step % 10 == 0:
            print('step: {}/{}... '.format(step, max_steps),
                  'loss: {:.4f}... '.format(batch_loss),
                  '{:.4f} sec/batch'.format((end - start)))
        if (step % 1000 == 0):
            saver.save(sess, os.path.join(save_path, 'model'), global_step=step)
        if step >= max_steps:
            break
    saver.save(sess, os.path.join(save_path, 'model'), global_step=step)
